import torch
#####################################################################
# Helper functions to make the code more readable.




def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)


# Compute log sum exp in a numerically stable way for the forward algorithm
def log_sum_exp(vec):
    # vec = [batch size, tagset_size]
    batch_size, tagset_size = vec.size()
    max_score, _ = torch.max(vec, 1)
    # max_score = [batch size]
    max_score_broadcast = max_score.view(batch_size, -1).expand(batch_size, tagset_size)
    return max_score + \
        torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))